package csci5828.tasks.impl;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.log4j.Logger;

import csci5828.tasks.RequestId;
import csci5828.tasks.TaskManager;
import csci5828.tasks.TaskResult;
import csci5828.tasks.TaskResults;
import csci5828.tasks.TaskSet;

class RequestContext<T> {
	
	private AtomicInteger done;
	private int total;
	private List<TaskResult<T>> results;
	private List<BlockingQueue<TaskResult<T>>> queues;
	private RequestId requestId;
	
	RequestContext(RequestId requestId, int total) {
		this.requestId = requestId;
		this.total = total;
		this.done = new AtomicInteger();
		this.results = Collections.synchronizedList(new ArrayList<TaskResult<T>>());
		this.queues = Collections.synchronizedList(new ArrayList<BlockingQueue<TaskResult<T>>>());
	}

	public RequestId getRequestId() {
		return requestId;
	}
	
	public int getDone() {
		return this.done.get();
	}
	
	public int getTotal() {
		return this.total;
	}
	
	public boolean isFinished()	{
   		return this.getDone() == this.getTotal();
	}
	
	public List<TaskResult<T>> getAvailableResults() {
		return Collections.unmodifiableList(this.results);
	}
	
	public void addResult(TaskResult<T> taskResult) {
		this.results.add(taskResult);
		this.done.incrementAndGet();
		for (BlockingQueue<TaskResult<T>> queue : this.queues) {
			queue.offer(taskResult);
		}
		if (this.done.get() == this.total) {
    		// To indicate that there are no more results, we add a null element 
    		for (BlockingQueue<TaskResult<T>> queue : this.queues) {
    			queue.offer(new TaskResultImpl<T>(null, getRequestId()));
    		}
		}
	}
	
	public BlockingQueue<TaskResult<T>> newQueue() {
		BlockingQueue<TaskResult<T>> queue = new LinkedBlockingQueue<TaskResult<T>>();
   		this.queues.add(queue);
		return queue;
	}
	
}

public class TaskManagerImpl<T> implements TaskManager<T> {
	
	private Logger logger = Logger.getLogger(getClass());
	private ExecutorService completionHandler;
	private CompletionService<TaskResult<T>> completionService;
	private ConcurrentHashMap<RequestId, RequestContext<T>> requestMap;
	private AtomicBoolean finished;
	
	public TaskManagerImpl() {
		ExecutorService executorService = Executors.newCachedThreadPool(new DaemonThreadFactory());
		this.completionService = new ExecutorCompletionService<TaskResult<T>>(executorService);
		this.completionHandler = Executors.newSingleThreadExecutor(new DaemonThreadFactory());
		this.requestMap = new ConcurrentHashMap<RequestId, RequestContext<T>>();
		this.finished = new AtomicBoolean();
		init();
	}
	
	private void init() {
		this.completionHandler.submit(new Runnable(){
			@Override
			public void run() {
				handleCompletedTasks();
			}
		});
	}

	@Override
	public RequestId submit(TaskSet<T> taskSet) {
		RequestId requestId = taskSet.getRequestId();
		if (requestMap.get(requestId) == null) {
			int size = taskSet.getTasks().size();
			RequestContext<T> requestContext = requestMap.putIfAbsent(taskSet.getRequestId(), new RequestContext<T>(requestId,size));
			if (requestContext != null)
				throw new RuntimeException("Concurrent submit for requestId="+requestId);
		} else {
			throw new RuntimeException("Already submitted; requestId="+requestId);
		}
		for (Callable<TaskResult<T>> callable : taskSet.getTasks()) {
			this.completionService.submit(callable);
		}
		return requestId;
	}

	@Override
	public boolean isFinished(RequestId requestId) {
		RequestContext<T> context = requestMap.get(requestId);
		if (context != null) {
			return context.isFinished();
		}
		return false;
	}

	private void handleCompletedTasks() {
		while (!finished.get()) {
    		try {
        		Future<TaskResult<T>> future = this.completionService.take();
        		TaskResult<T> taskResult = future.get();
        		logger.info("Got result for requestId="+taskResult.getRequestId());
        		RequestContext<T> requestContext = requestMap.get(taskResult.getRequestId());
        		if (requestContext != null) {
        			requestContext.addResult(taskResult);
        		} else {
        			System.err.println("requestContext == null");
        			assert(false);
        		}
    		} catch (InterruptedException e) {
    			e.printStackTrace();
    		} catch (ExecutionException e) {
    			e.printStackTrace();
    		}
		}
	}

	@Override
	public TaskResults<T> getResults(RequestId requestId) {
    	RequestContext<T> requestContext = requestMap.get(requestId);
    	if (requestContext != null) {
    		return new TaskResultsImpl<T>(
    			requestContext.getAvailableResults(), 
    			requestContext.newQueue(),
    			!requestContext.isFinished()
    		);
    	}
    	return null;
	}

}
